from textual.widgets import DataTable
from textual.events import MouseEvent
from screens.message import MessageScreen

import psutil
import datetime


class Users(DataTable):

    def __init__(self, id):
        super().__init__(id=id)
    
    def key_tab(self):
        self.app.query_one("#pid_table").styles.border = ("solid", "orangered")
        self.styles.border = ("solid", "white")
        self.app.query_one("#pid_table").cursor_type = 'row'
        self.cursor_type = 'none'

    def update_user_label(self):
        self.border_title = "在线用户:"+str(len(psutil.users()))

    def kill_process(self, pid):
        try:
            parent_proc = psutil.Process(pid)
            for child_proc in parent_proc.children(recursive=True):
                child_proc.kill()
            parent_proc.kill()
        except:
            pass

    def key_space(self):
        if self.has_focus:
            row = self.get_row_at(self.cursor_row)
            row_key = str(row[0])+str(row[1])+str(row[3])
            pid = row[3]
            for i in psutil.users():
                if str(i.name)+str(i.terminal)+str(i.pid)==row_key:
                    user = i
                    break
            user_start_time = datetime.datetime.fromtimestamp(user.started).strftime("%Y-%m-%d %H:%M:%S")
            self.app.install_screen(
                MessageScreen(
                    text=f"\n提示: 是否踢出用户{row[0]}?\n\n终端: {row[1]}  PID: {row[3]}\nIP: {row[2]}\n登录时间: {user_start_time}"), 
                    name="message")
            
            def check_result(result: bool):
                if result:
                    self.kill_process(pid)
                    self.remove_row(str(row_key))
                self.app.uninstall_screen("message")
            self.app.push_screen("message", check_result)

    def on_click(self, event: MouseEvent):
        # 1: left   2:middle   3:right
        if event.button==3 and self.has_focus:
            row = self.get_row_at(self.cursor_row)
            row_key = str(row[0])+str(row[1])+str(row[3])
            pid = row[3]
            for i in psutil.users():
                if str(i.name)+str(i.terminal)+str(i.pid)==row_key:
                    user = i
                    break
            user_start_time = datetime.datetime.fromtimestamp(user.started).strftime("%Y-%m-%d %H:%M:%S")
            
            self.app.install_screen(
                MessageScreen(
                    text=f"\n提示: 是否踢出用户{row[0]}?\n\n终端: {row[1]}  PID: {row[3]}\nIP: {row[2]}\n登录时间: {user_start_time}"), 
                    name="message")
            def check_result(result: bool):
                if result:
                    self.kill_process(pid)
                    self.remove_row(str(row_key))
                self.app.uninstall_screen("message")
            self.app.push_screen("message", check_result)
            
        if event.button==1 and self.has_focus:
            self.styles.border = ("solid", "orangered")
            self.app.query_one("#tree_box").styles.border = ("solid", "white")
            self.app.query_one("#pid_table").styles.border = ("solid", "white")
            self.cursor_type = 'row'
            self.app.query_one("#pid_table").cursor_type = 'none'

    def on_mount(self):
        self.header = ("用户名", "终端", "IP", "pid")
        self.cursor_type = 'none'
        self.zebra_stripes = True
        for j in self.header:
            self.add_column(j, key=j)
        self.users = psutil.users()
        for i in self.users:
            self.add_row(i.name, 
                         i.terminal,
                         '-' if '' == i.host else i.host,
                         i.pid,
                         key=str(i.name)+str(i.terminal)+str(i.pid))
        # 检测用户刷新时间2s
        self.update_timer = self.set_interval(2, self.update_users, pause=False)

    def update_users(self):
        self.update_user_label()
        row_keys = []
        users = psutil.users()
        row_keys = [str(i.name)+str(i.terminal)+str(i.pid) for i in users]
        # 删除已经下线的用户
        for i in range(len(self.rows)):
            row = self.get_row_at(i)
            row_key = str(row[0])+str(row[1])+str(row[3])
            if row_key not in row_keys:
                self.remove_row(str(row_key))
        # 添加新用户
        new_users = list(set(users)-set(self.users))
        for j in new_users:
            self.add_row( 
                j.name, 
                j.terminal, 
                '-' if '' == j.host else j.host,
                j.pid,
                key=str(j.name)+str(j.terminal)+str(j.pid))
        self.users = users


